from core.function import Function
from utils.common import to_variable


# 改变形状
class Reshape(Function):
    def __init__(self, shape):
        """
        初始化 Reshape 函数。

        Parameters:
            shape (tuple): 重塑后的形状。
        """
        self.x_shape = None
        self.shape = shape

    def forward(self, x):
        """
        执行Reshape操作的前向传播。

        Parameters:
            x (Variable): 输入变量。

        Returns:
            Variable: 重塑后的结果。
        """
        self.x_shape = x.shape
        y = x.reshape(self.shape)
        return y

    def backward(self, gy):
        """
        执行Reshape操作的反向传播。

        Parameters:
            gy (Variable): 相对于前向传播结果的梯度。

        Returns:
            Variable: 相对于输入变量的梯度。
        """
        # 重塑的反向传播将其转换成输入变量的形状
        return Reshape(self.x_shape).reshape(gy.data)

    def reshape(self, x):
        """
        执行Reshape操作。

        Parameters:
            x (Variable): 输入变量。

        Returns:
            Variable: 重塑后的结果。
        """
        if x.shape == self.shape:
            return to_variable(x)
        return self(x)



if __name__ == '__main__':
    from utils.functions_collect import reshape
    from core.variable import Variable

    x = Variable([[0, 1, 2], [3, 4, 5]])
    y = reshape(x, 6)  # y = x.reshape(6)
    y.backward(retain_grad=True)
    print(x.grad)
